from Shell.core.model import initialize_client, create_completion
from colorama import Fore, Style

from Shell.utils.config_manager import get_model_config

# 从配置文件中获取模型名称
model_config = get_model_config("Model")
model_prompt = model_config.get("prompt",
                                """- Role: PowerShell scripting expert and command line assistant
                                      - Background: Users need assistance when working with the Windows operating system, particularly in PowerShell scripting and command line operations.
                                      - Profile: You are an expert familiar with the Windows operating system and PowerShell, capable of providing concise and efficient solutions.
                                      - Skills: You possess profound skills in PowerShell script writing and command line operations, enabling you to quickly understand and execute user needs.
                                      - Goals: To provide a PowerShell command or solution that meets the specific needs of users in the Windows environment, including reading files with UTF-8 encoding.
                                      - Constrains: Must use Markdown code block syntax to display code, YAML, JSON, or XML, excluding tables. For PowerShell commands, use ```powershell to encapsulate. Command line commands must be fully displayed on the same line, without line breaks. Responses should be concise, without additional commands, providing only one command. For questions not executing commands, respond using ```powershell echo`. Include UTF-8 when reading files.
                                      - OutputFormat: Depending on user needs, the output format may be a PowerShell command, an explanation, or a simple echo command.
                                      - Workflow:
                                        1. Understand the user's needs and context.
                                        2. Provide a PowerShell command or solution based on the needs, ensuring UTF-8 encoding is used when reading files.
                                        3. Correctly encapsulate the command or output using Markdown code block syntax.
                                        4. 文件里写入多行内容的时候以powershell脚本来写入。
                                      - Examples:
                                        - Example 1: The user needs to list all files and folders in the current directory.
                                          ```powershell
                                          Get-ChildItem
                                          ```
                                        - Example 2: The user wants to know how to create a new folder.
                                          ```powershell
                                          New-Item -ItemType Directory -Name "NewFolder"
                                          ```
                                        - Example 3: The user needs to read a file with UTF-8 encoding.
                                          ```powershell
                                          Get-Content -Path "filename.txt" -Encoding UTF8
                                          ```
                                        - Example 4: When the user asks a question that requires an answer.
                                          ```powershell
                                          echo <answer content>
                                          ```
                              """)

conversation_history = [
    {
        "content": model_prompt,
        "role": "system",
    }
]


def clear_history():
    # 清理聊天历史
    global conversation_history
    conversation_history = [
        {
            "content": model_prompt,
            "role": "system",
        }
    ]
    print(Fore.BLUE + "Assistant: " + Style.RESET_ALL + "Chat history cleared.")


from mistralai.models.sdkerror import SDKError


def get_response(user_message, model_name="pixtral-large-latest"):
    # 添加用户消息到对话历史
    conversation_history.append({"content": user_message, "role": "user"})

    # 初始化模型客户端
    model_client = initialize_client(model_name)

    try:
        # 创建聊天补全请求
        res = create_completion(model_client, conversation_history, model_name)
    except SDKError as e:
        handle_sdk_error(e)
        # Return an appropriate message to the user
        return "Sorry, I'm unable to process your request right now. Please try again later."

    # 处理响应
    if res is not None:
        assistant_message = res.choices[0].message.content
        # 添加助手消息到对话历史
        conversation_history.append({"content": assistant_message, "role": "assistant"})
        return assistant_message
    else:
        return "No response from the model."


def handle_sdk_error(error):
    if error.status_code == 429:
        # Rate limit exceeded
        print("Rate limit exceeded. Please wait before making more requests.")
        # Optionally implement a retry mechanism with exponential backoff
    else:
        # Handle other SDK errors
        print(f"An SDK error occurred: {error}")
